#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Oct 22 14:45:16 2016

@author: kevin
"""
import tensorflow as tf
from utils import weight_variable, bias_variable

class GlimpseNet(object):
    """ Glimepse Network
    input  --- loction and alpha matte 
    output --- feature gt for rnn
    """
    def __init__(self, config,):
        self.alphaMatte_size = config.alphaMatte_sz             # alpha matte size  
        self.alphaMatte_channels = config.depth   # alpha matte channels == 1

        self.win_size = 75          # win for extracting glimpse  64
        self.depth = config.depth                # set to be 1
        self.sensor_size = self.win_size*self.win_size*self.depth    # sensor size 
        
        
        self.g_size = config.g_size      # glimpse size for l_t and alpha_t      256
        self.hg_size = config.hg_size    # hidden a-layer size                   1024
        self.hl_size = config.hl_size    # hidden l-layer size                   128
        self.loc_dim = config.loc_dim    # loc vector dim                        2
        
        self.init_weights()
        
    def init_weights(self):
        """ initialize trainable weights for l_t
            loc_dim -> hl_size -> g_size 
        """
        # x(1,2) * W(2,128) + b(1,128)
        self.w_l0 = weight_variable((self.loc_dim-1, self.hl_size))               
        self.b_l0 = bias_variable((self.hl_size,))
        # x(1,128) * W(128,256) + b(1,256)
        self.w_l1 = weight_variable((self.hl_size, self.g_size))
        self.b_l1 = bias_variable((self.g_size,))
        """ initialize trainable weights for alpha_t """
        """ sensor_size -> hg_size -> g_size """
        # x(1,4096) * W(4096,1024) + b(1,1024)
        self.w_g0 = weight_variable((self.sensor_size, self.hg_size))
        self.b_g0 = bias_variable((self.hg_size,))
        # x(1,1024) * W(1024,256) + b(1,256)
        self.w_g1 = weight_variable((self.hg_size, self.g_size))
        self.b_g1 = bias_variable((self.g_size,))
        
    def get_glimpse(self, loc):
        """take glimpse on the alpha matte"""
        #aMattes = tf.reshape(self.alpha_ph, [-1,self.alphaMatte_size, self.alphaMatte_size,self.alphaMatte_channels])
        aMattes = self.alpha_ph

#################3 layers cnn, self.alphaMatte_size is fixed as 600
#        weight_shape=[5,5,self.alphaMatte_channels,self.alphaMatte_channels]
#        w1=weight_variable(weight_shape)
#        b1=bias_variable([self.alphaMatte_channels])
#        conv1=tf.nn.conv2d(aMattes,w1,strides=[1,2,2,1],padding='SAME')+b1   #[-1,600,600,1]==>[-1,300,300,1]
#
#        w2=weight_variable(weight_shape)
#        b2=bias_variable([self.alphaMatte_channels])
#        conv2=tf.nn.conv2d(conv1,w2,strides=[1,2,2,1],padding='SAME')+b2   #[-1,300,300,1]==>[-1,150,150,1]
#
#        w3=weight_variable(weight_shape)
#        b3=bias_variable([self.alphaMatte_channels])
#        conv3=tf.nn.conv2d(conv2,w3,strides=[1,2,2,1],padding='SAME')+b3   #[-1,150,150,1]==>[-1,75,75,1]
#
#        glimpse_aMattes = tf.reshape(conv3, [-1,self.win_size*self.win_size*self.alphaMatte_channels])
 

      
######################################################################
        #glimpseNet by Xuke
        glimpse_aMattes = tf.image.extract_glimpse(aMattes, [self.win_size, self.win_size], loc) 
                        # result is 4D tensor [batch_size, glimpse_height, glimpse_width, channels]
        glimpse_aMattes = tf.reshape(glimpse_aMattes, [-1,self.win_size*self.win_size*self.alphaMatte_channels]) 
                        # feed this 2D tensor to the glimpse network
        #########################################################   

     

        return glimpse_aMattes   
        
    def __call__(self, loc, alpha_est):
        self.alpha_ph = alpha_est         # shape = [M, 512*512]
        glimpse_input = self.get_glimpse(loc)
        g = tf.nn.relu(tf.nn.xw_plus_b(glimpse_input, self.w_g0, self.b_g0))  
        g = tf.nn.xw_plus_b(g, self.w_g1, self.b_g1)                          
        l = tf.nn.relu(tf.nn.xw_plus_b(tf.cast(loc,tf.float64), tf.cast(self.w_l0,tf.float64), tf.cast(self.b_l0,tf.float64)))            
        l = tf.nn.xw_plus_b(tf.cast(l,tf.float64), tf.cast(self.w_l1,tf.float64), tf.cast(self.b_l1,tf.float64))  
        g_l=tf.cast(g,tf.float64)+l                        
        g = tf.nn.relu(g_l)
        g=tf.cast(g,tf.float32)                                                 
        return g   # this is a 2D tensor [M*batch_sz, g_size] #
        
class LocNet(object):
        """ Loc Network
        input  --- output feature vector ht of LSTM
        output --- next location and sampled next location
        
        loc is a 2D tensor [M*1, 2] , batch_sz = 1
        """
        def __init__(self, config):
            self.loc_dim = config.loc_dim              # loc dim (1,2)
            self.input_dim = config.rnn_output_sz      # RNN output size   (1,1000) == VGG output_sz
            self.loc_std = config.loc_std              # stddev for location sampling
            self._sampling = True
            self.r=config.win_img
            self.init_weights()
        
        def init_weights(self):
            # x(1,1000) * W(1000, 2) + b(1,2)
            self.w = weight_variable((self.input_dim, self.loc_dim))
            self.b = bias_variable((self.loc_dim,))


        
        def __call__(self, input):
            mean=tf.nn.xw_plus_b(tf.cast(input,tf.float64), tf.cast(self.w,tf.float64), tf.cast(self.b,tf.float64)) 
            mean = tf.stop_gradient(mean) # stop gradient computation
            if self._sampling:
                loc = mean + tf.random_normal((1, self.loc_dim), stddev=self.loc_std,dtype=tf.float64)
            else:
                loc =mean
            loc = tf.stop_gradient(loc)
            loc=tf.cast(loc,tf.float32)
#####################2017-3-9#################
            loc_loc=tf.clip_by_value(loc[:,0:2],-1,1)
            loc_loc=tf.round(loc_loc/self.r)*self.r
            loc_ratio=tf.clip_by_value(loc[:,2],0.8,1.2)
##################################################
            return loc,mean,loc_loc,loc_ratio            # loc is the sampled by mean
        
        @property
        def sampling(self):
            return self._sampling
            
        @sampling.setter
        def sampling(self, sampling):
            self._sampling = sampling


